use opencv::core::{Mat, MatTrait, CV_8UC4};
use windows::{
  core::*, Win32::Graphics::Direct3D::*, Win32::Graphics::Direct3D11::*, Win32::Graphics::Dxgi::*,
};

pub struct Capturer {
  duplicate_output: IDXGIOutputDuplication,
  device: ID3D11Device,
  context: ID3D11DeviceContext,
  pub width: i32,
  pub height: i32,
}

impl Capturer {
  pub fn new() -> Result<Capturer> {
    let (device, context) = create_device()?;
    let dxdevice = device.cast::<IDXGIDevice>()?;
    let adapter = unsafe { dxdevice.GetAdapter()? };
    let output = unsafe { adapter.EnumOutputs(0)? };
    let output1 = output.cast::<IDXGIOutput1>()?;
    let duplicate_output = unsafe { output1.DuplicateOutput(&device)? };
    let mut desc = DXGI_OUTPUT_DESC::default();
    unsafe { output.GetDesc(&mut desc) }?;
    let width = desc.DesktopCoordinates.right - desc.DesktopCoordinates.left;
    let height = desc.DesktopCoordinates.bottom - desc.DesktopCoordinates.top;
    Ok(Capturer {
      duplicate_output,
      device,
      context,
      width,
      height,
    })
  }

  pub fn capture_frame(&self) -> Result<Mat> {
    let mut pframeinfo = DXGI_OUTDUPL_FRAME_INFO::default();
    let mut ppdesktopresource = None;
    unsafe {
      self
        .duplicate_output
        .AcquireNextFrame(1000, &mut pframeinfo, &mut ppdesktopresource)?;
      if ppdesktopresource.is_none() {
        return Err(Error::from_win32());
      }
      let source = ppdesktopresource.unwrap();
      let acquire_frame = source.cast::<ID3D11Texture2D>()?;
      let mut desc = D3D11_TEXTURE2D_DESC::default();
      acquire_frame.GetDesc(&mut desc);

      desc.Usage = D3D11_USAGE_STAGING;
      desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ.0 as u32;
      desc.BindFlags = 0;
      desc.MiscFlags = 0;
      desc.MipLevels = 1;
      desc.ArraySize = 1;
      desc.SampleDesc.Count = 1;
      desc.SampleDesc.Quality = 0;
      let mut texture: Option<ID3D11Texture2D> = None;

      self
        .device
        .CreateTexture2D(&desc, None, Some(&mut texture))?;

      let _texture = texture.unwrap();
      self.context.CopyResource(&_texture, &acquire_frame);

      let surface = _texture.cast::<IDXGISurface1>()?;
      let mut map = DXGI_MAPPED_RECT::default();
      surface.Map(&mut map, DXGI_MAP_READ)?;
      self.duplicate_output.ReleaseFrame()?;

      // println!("{} {} {}", map.Pitch, desc.Width, desc.Height);
      let w = desc.Width;
      let h: u32 = desc.Height;
      if let Ok(mut img) = Mat::new_nd(&[h as i32, w as i32], CV_8UC4) {
        map.pBits.copy_to(img.data_mut(), (w * h * 4) as usize);
        return Ok(img);
      }
      Err(Error::from_win32())
    }
  }
}

fn create_device() -> Result<(ID3D11Device, ID3D11DeviceContext)> {
  let _type = D3D_DRIVER_TYPE_HARDWARE;
  let mut device = None;
  let mut context = None;
  let flags = D3D11_CREATE_DEVICE_BGRA_SUPPORT;

  let feature_levels = [
    D3D_FEATURE_LEVEL_11_0,
    D3D_FEATURE_LEVEL_10_1,
    D3D_FEATURE_LEVEL_10_0,
    D3D_FEATURE_LEVEL_9_1,
  ];

  unsafe {
    D3D11CreateDevice(
      None,
      _type,
      None,
      flags,
      Some(&feature_levels),
      D3D11_SDK_VERSION,
      Some(&mut device),
      None,
      Some(&mut context),
    )?;
  }
  Ok((device.unwrap(), context.unwrap()))
}
